from typing import Union

import yaml
import fire
import tqdm
import os
import sys
from pathlib import Path

sys.path.append(str(Path(__file__).absolute().parent))

import torch
import torchaudio
import torch.multiprocessing as mp

from model.build_model import build_model
from utils.load_scp import read_2column_text

def normalize(raw_wav, est_wav):
    norm = torch.norm(raw_wav, p=float('inf'))
    est_wav = est_wav * norm / torch.max(torch.abs(est_wav))
    return est_wav


def run(rank, world_size, config, ckpt, mix_scp, ref_scp, output_dir, device):
    device = device if world_size == 1 else device[rank]
    with open(config, "r") as f:
        config = yaml.safe_load(f)
    ckpt = torch.load(ckpt, map_location='cpu')
    model = build_model(config['model_conf'])
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device)
    mix_scp_dict = read_2column_text(mix_scp)
    ref_scp_dict = read_2column_text(ref_scp)
    keys = list(mix_scp_dict.keys())
    keys = keys[rank::world_size]
    with torch.no_grad():
        for _k in tqdm.tqdm(keys, desc=f"[rank {rank}]"):
            mix_path, ref_path = mix_scp_dict[_k], ref_scp_dict[_k]
            save_path = str(Path(output_dir) / f"{Path(mix_path).stem}.wav")
            mix_audio, sr = torchaudio.load(mix_path)
            ref_audio, sr = torchaudio.load(ref_path)
            output = model(mix_audio.to(device), ref_audio.to(device))
            output = normalize(mix_audio, output)
            torchaudio.save(save_path, output.cpu(), sample_rate=sr)
    pass


def main(config, ckpt, mix_scp, ref_scp, output_dir, device: Union[list, str]):
    """
    
    """
    os.makedirs(output_dir, exist_ok=True)
    if isinstance(device, list):
        print(f"running {len(device)} processes")
        mp.spawn(run, nprocs = len(device), args=(len(device), config, ckpt, mix_scp, ref_scp, output_dir, device))
    else:
        run(0, 1, config, ckpt, mix_scp, ref_scp, output_dir, device)
    print("Done...")


if __name__ == "__main__":
    fire.Fire(main)
    pass